#!/usr/bin/env python
# -*- coding:utf-8 -*-
#
#   Author  :   eleven.i386
#   E-mail  :   eleven.i386@gmail.com
#   Date    :   13/03/08 12:02:28
#   Desc    :   traffic_core 流量采集模块; update 内容更新模块; other_win,main; 主窗体
#   Use     :   需要sudo执行,或者以root用户执行.开启网卡混杂模式需要root
#

import gtk
import getpass
import Queue
import Image
import ImageDraw
import ConfigParser
import os,sys,re
import gobject
import thread
import threading
import pcap,dpkt,socket
from time import sleep,ctime
from subprocess import Popen, PIPE

class traffic_core():
    """ 利用pcap进行流量收集"""
    def __init__(self):
        self.trafficQ = []
        self.now_proc = {}
        self.Flag = False
        self.total = self.tmp = ""
        self.pre = re.compile("^\-")
        self.total_traffic = {}

    def get_data(self):
        Pcap = pcap.pcap(self.get_Iface())
        Pcap.setfilter("tcp")
        sip = dip = sport = dport = ""
        net_info = {}
        second = 0
        now_net = {}
        process = self.get_pid()

        for ptime, pdata in Pcap:
            item = dpkt.ethernet.Ethernet(pdata)
            sip = "%d.%d.%d.%d"%(tuple(map(ord,list(item.data.src))))
            dip = "%d.%d.%d.%d"%(tuple(map(ord,list(item.data.dst))))
            if item.data.data.__class__.__name__ == 'TCP':
                sport = item.data.data.sport
                dport = item.data.data.dport
                data_len = len(item.data.data.data)
            if second == 0:second = ptime
            if float(ptime) - float(second) < 1.0:
                key = "%s:%s %s:%s"%(sip,sport,dip,dport)
                if net_info.has_key(key):net_info[key] = net_info[key] + data_len
                else:net_info[key] = data_len
            else:second = 0;break

        for key in net_info:
            if process.has_key(key):
                sip = key.split()[0].split(':')[0]
                if sip == self.get_Ip():
                    Up = round(float(net_info[key])/1024,2)
                    Down = 0.00
                    if now_net.has_key(process[key]):
                        d,u = float(now_net[process[key]][0])+Down,float(now_net[process[key]][1])+Up
                        now_net[process[key]] = d,u
                    else:
                        now_net[process[key]] = Down,Up
                else:
                    Up = 0.00
                    Down = round(float(net_info[key])/1024,2)
                    if now_net.has_key(process[key]):
                        d,u = float(now_net[process[key]][0])+Down,float(now_net[process[key]][1])+Up
                        now_net[process[key]] = d,u
                    else:
                        now_net[process[key]] = Down,Up
        self.trafficQ.append(now_net)

    def get_pid(self):
        d_net_info = {}
        n = os.popen('netstat -tnop').readlines()
        for i in n[2:]:
            d_net_info[i.split()[3]+" "+i.split()[4]] = i.split()[6]
            d_net_info[i.split()[4]+" "+i.split()[3]] = i.split()[6]
        return d_net_info

    def get_Iface(self):
        return  ''.join([ i.split()[-1] for i in os.popen('netstat -rn').readlines()[2:] if i.split()[0] == '0.0.0.0' ])

    def get_Ip(self):
        return re.search('\d+\.\d+\.\d+\.\d+',Popen('ifconfig', stdout=PIPE).stdout.read()).group(0)

    def process_view(self):
        buf = []
        proc_dict = {}
        process = set(self.get_pid().values())
        for i in process:
            if not self.pre.match(i): proc_dict[i] = (0.00,0.00)
        for i in proc_dict.keys():
            if i not in self.total_traffic.keys():self.total_traffic[i] = proc_dict[i]
        if len(self.trafficQ) == 0:
            if self.Flag:
                proc_dict.update(self.now_proc)
                self.Flag = True
                for i in proc_dict:
                    if not self.pre.match(i):
                        buf.append((int(i.split('/')[0]),i.split('/')[1],str(proc_dict[i][0])+' KB/s', str(proc_dict[i][1])+' KB/s',\
                            str(self.total_traffic[i][0])+' KB', str(self.total_traffic[i][1])+' KB'))
            else:
                for i in proc_dict:
                    if not self.pre.match(i):
                        self.total_traffic[i] = (self.total_traffic[i][0] + proc_dict[i][0],self.total_traffic[i][1] + proc_dict[i][1])
                        buf.append((int(i.split('/')[0]),i.split('/')[1],str(proc_dict[i][0])+' KB/s', str(proc_dict[i][1])+' KB/s',\
                           str(self.total_traffic[i][0])+' KB', str(self.total_traffic[i][1])+' KB'))
                self.Flag = False
            self.now_proc = {}
        else:
            self.now_proc = self.trafficQ.pop()
            proc_dict.update(self.now_proc)
            for i in proc_dict:
                if not self.pre.match(i):
                    self.total_traffic[i] = (self.total_traffic[i][0] + proc_dict[i][0],self.total_traffic[i][1] + proc_dict[i][1])
                    buf.append((int(i.split('/')[0]),i.split('/')[1],str(proc_dict[i][0])+' KB/s', str(proc_dict[i][1])+' KB/s',\
                        str(self.total_traffic[i][0])+' KB', str(self.total_traffic[i][1])+' KB'))
            self.Flag = False
        return buf

    def total_view(self):
        up, down = 0.00,0.00
        if len(self.now_proc) == 0:
            if self.Flag:
                self.Flag = False
                return self.total
            else:
                self.total = 'R: %s KB/S | S: %s KB/S'%(down, up)
        else:
            for i in self.now_proc:
                up += self.now_proc[i][1]
                down += self.now_proc[i][0]
            self.total = 'R: %s KB/S | S: %s KB/S'%(down, up)
        return self.total

    def result(self):
        proc_dict = {}
        process = set(self.get_pid().values())
        for i in process:
            if not self.pre.match(i): proc_dict[i] = (0.00,0.00)
        for i in proc_dict.keys():
            if i not in self.total_traffic.keys():self.total_traffic[i] = proc_dict[i]
        if len(self.tmp) == 0 or len(self.tmp[0]) == 0:
            r = []
            for i in proc_dict:
                if not self.pre.match(i):
                    r.append((int(i.split('/')[0]),i.split('/')[1],str(proc_dict[i][0])+' KB/s', str(proc_dict[i][1])+' KB/s',\
                       str(self.total_traffic[i][0])+' KB', str(self.total_traffic[i][1])+' KB'))
            return (r,"R: 0.00 KB/S S :0.00 KB/S")
        else:return self.tmp

    def data_loop(self):
        while 1:
            self.get_data()

    def result_loop(self):
        while 1:
            self.tmp = (self.process_view(),self.total_view())
            sleep(1)

class update():
    """ 更新主界面 """
    def __init__(self):
        self.lstore = self.img =  None
        self.p_some, self.t_some = [],[]
        self.A = traffic_core()
        thread.start_new_thread(self.A.data_loop,())
        thread.start_new_thread(self.A.result_loop,())

    def setImage(self,d):
        self.im = Image.new("RGB",(180,22),(30,144,255))
        self.dr = ImageDraw.Draw(self.im)
        self.dr.text((5,5),d,fill="#FFFFFF")
        self.im.save("/tmp/.t.png")
        self.img_buf = gtk.gdk.pixbuf_new_from_file("/tmp/.t.png")
        self.img.set_from_pixbuf(self.img_buf)

    def __add_data(self,lstore):
        self.lstore.clear()

        for item in lstore:
            self.lstore.append([item[0], item[1], item[2], item[3], item[4], item[5], "正在运行"])
        return self.lstore

    def Process_callback(self):
        if self.p_some == self.A.result()[0]:return True
        self.__add_data(self.A.result()[0])
        self.p_some = self.A.result()[0][:]
        return True

    def Image_callback(self):
        if self.t_some == self.A.result()[1]:return True
        self.setImage(self.A.result()[1])
        self.t_some = self.A.result()[1][:]
        return True

class other_win(gtk.Window):
    """ 显示进程的流量的列表 """
    def __init__(self):
        super(other_win,self).__init__()
        # 取消边框
        self.set_decorated(False)
        # 不在任务栏显示
        self.set_skip_taskbar_hint(True)

        # 一个窗口 用来显示详细
        self.sw = gtk.ScrolledWindow()
        self.sw.set_shadow_type(gtk.SHADOW_ETCHED_IN)
        self.sw.set_policy(gtk.POLICY_NEVER, gtk.POLICY_AUTOMATIC)
        self.set_default_size(300,300)
        self.u = update()
        self.u.lstore = self.__create_model()
        treeview = gtk.TreeView(self.u.lstore)
        treeview.set_rules_hint(True)
        self.sw.add(treeview)
        self.__add_columns(treeview)
        self.vbox = gtk.VBox()
        self.vbox.pack_start(self.sw)
        self.add(self.vbox)
        self.vbox.show()
        self.sw.show()

    def __create_model(self):
        lstore = gtk.ListStore(int, str, str, str, str, str, str )
        return lstore

    def __add_columns(self, treeview):
        rendertext = gtk.CellRendererText()

        column = gtk.TreeViewColumn('PID', rendertext, text=0)
        column.set_sort_column_id(0)
        treeview.append_column(column)

        column = gtk.TreeViewColumn('进程名称', rendertext, text=1)
        column.set_sort_column_id(1)
        treeview.append_column(column)

        column = gtk.TreeViewColumn('下载速度', rendertext, text=2)
        column.set_sort_column_id(2)
        treeview.append_column(column)

        column = gtk.TreeViewColumn('上传速度', rendertext, text=3)
        column.set_sort_column_id(3)
        treeview.append_column(column)

        column = gtk.TreeViewColumn('下载流量', rendertext, text=4)
        column.set_sort_column_id(4)
        treeview.append_column(column)

        column = gtk.TreeViewColumn('上传流量', rendertext, text=5)
        column.set_sort_column_id(5)
        treeview.append_column(column)

        column = gtk.TreeViewColumn('进程状态', rendertext, text=6)
        column.set_sort_column_id(6)
        treeview.append_column(column)

        # 回調刷新窗體
        gobject.idle_add(self.u.Process_callback)

class main(gtk.Window):
    """ 主窗口,悬浮窗 """
    def __init__(self):
        super(main,self).__init__()
        self.flag = True
        # 只有按下左键时拖动窗体
        self.drag = False
        self.mouse_x, self.mouse_y = 0,0
        # 窗口置顶
        self.set_keep_above(True)
        # 主窗体初始大小
        self.set_default_size(230,22)
        # 初始坐标
        self.x, self.y = 1110, 28
        self.move(self.x, self.y)
        #self.set_opacity(1)
        # 主窗体响应鼠标点击事件
        self.add_events(gtk.gdk.BUTTON_PRESS_MASK|\
                        gtk.gdk.BUTTON_RELEASE_MASK|\
                        gtk.gdk.POINTER_MOTION_MASK|\
                        gtk.gdk.POINTER_MOTION_HINT_MASK)


        # 禁止调整窗口大小
        self.set_resizable(False)
        # 不在任务栏显示
        self.set_skip_taskbar_hint(True)
        # 取消边框
        self.set_decorated(False)

        self.ow = other_win()
        self.u = update()


        # 一个图片用来显示total
        self.u.img = gtk.Image()
        vbox = gtk.VBox()
        vbox.add(self.u.img)
        self.add(vbox)

        self.connect("button-press-event",self.mouse_click)
        self.connect("button-release-event",self.mouse_release)
        self.connect("motion-notify-event",self.mouse_move)
        self.show_all()

        # 回調刷新主窗體
        gobject.idle_add(self.u.Image_callback)

    # 鼠标移动事件
    def mouse_move(self,widget,event,data=None):
        if self.drag:
            x,y = self.get_position()
            self.move((x+int(event.x)-int(self.mouse_x)),\
                        (y+int(event.y)-int(self.mouse_y)))

    # 点击释放事件
    def mouse_release(self,widget,event,data=None):
        if event.button == 1:self.drag = False

    # 点击按下事件
    def mouse_click(self,widget,event,data=None):
        if event.type == gtk.gdk._2BUTTON_PRESS:
            if self.flag:
                self.ow.show_all()
                self.flag = False
            else:
                self.ow.hide()
                self.flag = True
            return

        if event.button == 3:
            self.destroy_quit(None, None)

        elif event.button == 1:
            self.drag = True
            self.mouse_x,self.mouse_y = event.x,event.y

    def destroy_quit(self, event, data=None):
        gtk.main_quit()

    def star(self):
        gtk.main()

if __name__ == "__main__":
    s = main()
    s.star()

